import multiprocessing
import threading
from collections import deque
from typing import Any, Callable, Iterable, Optional, Union

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Button

from .repeating_timer import RepeatingTimer
from .utils import convert_to_iterator

HistPoint = Union[float, Iterable[float]]
HistBufferData = Optional[tuple[np.ndarray, ...]]
HistViewData = Optional[tuple[tuple[np.ndarray, np.ndarray], ...]]


def plot_summary(data: HistViewData, title: Optional[str]):
    ys_histogram = data

    plt.figure(figsize=(8, 4))
    axes: plt.Axes = plt.gca()
    axes.grid()

    if title is not None:
        axes.set_title(title)

    # draw bars (and calculate min max)
    max_height = 0
    min_x, max_x = None, None
    for yi, (y_heights, y_edges) in enumerate(ys_histogram):
        axes.bar(y_edges[:-1], y_heights, y_edges[1:] - y_edges[:-1],
            align='edge', edgecolor="black", linewidth=1, zorder=2 + (1 - yi / len(ys_histogram)))
        max_height = max(max_height, np.max(y_heights))
        min_x = min(min_x, min(y_edges)) if min_x is not None else min(y_edges)
        max_x = max(max_x, max(y_edges)) if max_x is not None else max(y_edges)

    # adjust xy range
    padding = (max_x - min_x) / 32
    axes.set_xlim(min_x - padding, max_x + padding)
    axes.set_ylim(0, max_height * 1.1)

    plt.show()


class AccumulateHistBuffer:
    def __init__(self, size: int | None = 1024):
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)

        self._version = 0
        self._ys = deque[np.ndarray](maxlen=size)

    def append(self, ys: HistPoint):
        ys = np.asarray(ys, dtype=np.float32)
        if ys.ndim == 0:
            ys = ys.reshape([1])
        elif ys.ndim == 1:
            ys = ys.reshape([-1])
        else:
            raise ValueError(f"Expected ys.ndim = 0 or 1, got ys.shape={ys.shape}")

        with self._lock:
            self._version += 1
            self._ys.append(ys)
            self._condition.notify_all()

    def clear(self):
        with self._lock:
            self._version += 1
            self._ys.clear()
            self._condition.notify_all()

    def read(self) -> tuple[int, HistBufferData]:
        with self._lock:
            if len(self._ys) == 0:
                return self._version, None
            return self._version, tuple(self._ys)

    def wait_and_read(self, version: int) -> tuple[int, HistBufferData]:
        with self._lock:
            while version == self._version:
                self._condition.wait()
            return self.read()


def convert_to_view_data(data: HistBufferData, bins: int) -> HistViewData:
    if data is None:
        return None
    ys = np.transpose(data)  # [ys_n, seq_n]
    return tuple(np.histogram(y, bins=bins) for y in ys)


class AccumulateHistView:
    def __init__(self,
        figure: plt.Figure, *,
        title: Optional[str] = None,
    ):
        self._figure = self._init_figure(figure)
        self._axes = self._init_axes(figure, title)
        self._play_button, self._clear_button = self._init_buttons(figure)
        self._play_button.on_clicked(lambda _: self._on_play_clicked())
        self._clear_button.on_clicked(lambda _: self._on_clear_clicked())

        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

        self._paused = True
        self.request_pause: Optional[Callable[[], Any]] = None
        self.request_resume: Optional[Callable[[], Any]] = None
        self.request_clear: Optional[Callable[[], Any]] = None

    @staticmethod
    def _init_figure(figure: plt.Figure):
        figure.set_size_inches(8, 4)
        return figure

    @staticmethod
    def _init_axes(figure: plt.Figure, title: Optional[str] = None) -> plt.Axes:
        axes = figure.gca()
        axes.grid()

        if title is not None:
            axes.set_title(title)

        return axes

    @staticmethod
    def _init_buttons(figure: plt.Figure) -> tuple[plt.Button, plt.Button]:
        figure.subplots_adjust(0.1, 0.1, 0.95, 0.85)
        play_button = Button(figure.add_axes([0.01, 0.9, 0.09, 0.09]), "resume")
        clear_button = Button(figure.add_axes([0.11, 0.9, 0.09, 0.09]), "clear")
        return play_button, clear_button

    @staticmethod
    def _init_bars(axes: plt.Axes, ys_n: int, bars_n: int) -> list[tuple[plt.Rectangle, ...]]:
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def update(self, data: HistViewData):
        if data is None:  # clear
            self._axes.cla()
            self._axes.grid()
            self._ys_bars = None
            self._figure.canvas.draw_idle()
            return

        ys_histogram = data
        if self._ys_bars is None:
            ys_n = len(ys_histogram)
            if ys_n == 0:
                return
            bars_n = len(ys_histogram[0][0])
            if bars_n == 0:
                return
            self._ys_bars = self._init_bars(self._axes, ys_n, bars_n)

        max_height = 0
        min_x, max_x = None, None
        for y_bars, (y_heights, y_edges) in zip(self._ys_bars, ys_histogram):
            y_widths = y_edges[1:] - y_edges[:-1]
            for bar, x, width, height in zip(y_bars, y_edges[:-1], y_widths, y_heights):
                bar.set_x(x)
                bar.set_width(width)
                bar.set_height(height)
            max_height = max(max_height, np.max(y_heights))
            min_x = min(min_x, min(y_edges)) if min_x is not None else min(y_edges)
            max_x = max(max_x, max(y_edges)) if max_x is not None else max(y_edges)

        # adjust y range
        self._axes.set_ylim(0, max_height * 1.1)

        # adjust both x range
        padding = (max_x - min_x) / 32
        self._axes.set_xlim(min_x - padding, max_x + padding)

        # redraw
        self._figure.canvas.draw_idle()

    def pause(self):
        if self._paused:
            return
        self._paused = True
        self._play_button.label.set_text("resume")

    def resume(self):
        if not self._paused:
            return
        self._paused = False
        self._play_button.label.set_text("pause")

    def _on_play_clicked(self):
        if self._paused:
            if self.request_resume:
                self.request_resume()
        else:
            if self.request_pause:
                self.request_pause()

    def _on_clear_clicked(self):
        if self.request_clear:
            self.request_clear()


class AccumulateHistPopup:
    def __init__(self,
        producer: Union[Iterable[HistPoint], Callable[[], HistPoint]], *,
        produce_interval: float = 1 / 128,
        frame_interval: float = 1 / 32,
        bins: int = 20,
        title: Optional[str] = None,
        start: bool = True,
        buffer_size: int | None = 1024,
    ):
        """ Show histograms of accumulated data from `producer` in a popup window

        :param producer: an iterator (generator) or callable that produces data point with shape (ys_n,)
        :param produce_interval: interval between calls of `producer`
        :param frame_interval: interval between gui updates
        :param bins: number of bins in the histogram
        :param title: the title of the plot
        :param start: whether to start immediately
        :param buffer_size: the number of buffered data points. if None, the buffer is unbounded. 
        """
        self._iterator = convert_to_iterator(producer)
        self._data_connection, sub_data_connection = multiprocessing.Pipe()
        self._event_connection, sub_event_connection = multiprocessing.Pipe()
        self._subprocess = multiprocessing.Process(
            target=self.subprocess_func,
            args=(sub_data_connection, sub_event_connection,
                  frame_interval, bins, title, buffer_size))
        self._subprocess.start()
        self._buffer = AccumulateHistBuffer(buffer_size)

        def tick_func():
            point = next(self._iterator)
            try:
                self._data_connection.send(point)
            except (OSError, BrokenPipeError):
                self._timer.stop()
                return
            self._buffer.append(point)

        self._timer = RepeatingTimer(produce_interval, tick_func)

        def event_thread_func():
            while True:
                try:
                    cmd = self._event_connection.recv()
                except EOFError:
                    break
                if cmd == "pause":
                    self.pause()
                elif cmd == "resume":
                    self.resume()
                elif cmd == "clear":
                    self.clear()

        self._event_thread = threading.Thread(target=event_thread_func)
        self._event_thread.start()

        self._bins = bins
        self._title = title

        if start:
            self.resume()

    @staticmethod
    def subprocess_func(
        data_connection,
        event_connection,
        frame_interval: float,
        bins: int,
        title: Optional[str],
        buffer_size: Optional[int]
    ):
        matplotlib.use('QtAgg')

        buffer = AccumulateHistBuffer(buffer_size)

        def data_thread_func():
            while True:
                point = data_connection.recv()
                buffer.append(point)

        data_thread = threading.Thread(target=data_thread_func, daemon=True)
        data_thread.start()

        figure = plt.figure()
        view = AccumulateHistView(figure, title=title)
        view_version = -1

        def tick_func():
            nonlocal view_version
            view_version, data = buffer.wait_and_read(view_version)
            view_data = convert_to_view_data(data, bins)
            view.update(view_data)

        timer = RepeatingTimer(frame_interval, tick_func)
        timer.start()

        view.request_pause = lambda: event_connection.send("pause")
        view.request_resume = lambda: event_connection.send("resume")
        view.request_clear = lambda: event_connection.send("clear")

        def event_thread_func():
            while True:
                cmd = event_connection.recv()
                if cmd == "pause":
                    view.pause()
                elif cmd == "resume":
                    view.resume()
                elif cmd == "clear":
                    buffer.clear()

        event_thread = threading.Thread(target=event_thread_func, daemon=True)
        event_thread.start()

        plt.show()

    def resume(self):
        self._timer.start()
        self._event_connection.send("resume")

    def pause(self):
        self._timer.stop()
        self._event_connection.send("pause")

    def clear(self):
        self._buffer.clear()
        self._event_connection.send("clear")

    def wait(self):
        if self._subprocess.is_alive():
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def close(self):
        if self._subprocess.is_alive():
            self._subprocess.terminate()
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def plot_summary(self):
        _, data = self._buffer.read()
        view_data = convert_to_view_data(data, self._bins)
        plot_summary(view_data, title=self._title)

    def wait_and_plot_summary(self):
        self.wait()
        self.plot_summary()

    def close_and_plot_summary(self):
        self.close()
        self.plot_summary()
